# Geospatial analysis - Regional estimates of AA-linked signatures in Romanian and Serbian counties
rm(list = ls())

# Spatial objects
library(geojsonio)
library(sp)
library(sf)
library(raster)
# Data handling
library(dplyr)
library(stringr)
library(tidyr)
library(matrixStats)
library(Matrix)
# Mixed models
library(lme4)
library(tibble)
library(broom)
library(broom.mixed)
# Plots
library(ggplot2)
library(patchwork)

# Loading individual data
# Residency and duration for all cases along with signature attribution
aa_extended <- readRDS("datasets/aa_extended")
aa_corrected <- readRDS("datasets/aa_corrected")

# Loading maps
# Romanian counties - Grouping Ilfov & Bucuresti
ro_regions <- geojson_read("maps/stanford-rf707cg1788-geojson.json", what = "sp") %>% 
  aggregate(by = "nam")
ro_regions@data <- ro_regions@data %>% mutate(varname_1 = ifelse(nam %in% c("Municipiul Bucuresti", "Ilfov"), "Bucuresti/Ilfov", nam))
ro_regions <- aggregate(ro_regions, by="varname_1")
# Serbian counties
sr_regions <- geojson_read("maps/stanford-hf794fx7640-geojson.json", what = "sp")
sro_regions <- bind(sr_regions, ro_regions)

# Creating matrix with weights for each residencies and AA signatures
aa_matrix <- aa_extended %>% 
  mutate(region = case_when(is.na(region) ~ "Abroad",
                            region %in% sr_regions@data$varname_1 ~ str_c("SR:",region),
                            region %in% ro_regions@data$varname_1 ~ str_c("RO:",region))) %>% 
  filter(!is.na(region)) %>% 
  select(donor_id,region,dur) %>%
  group_by(donor_id, region) %>% 
  summarize(dur=sum(dur)) %>% 
  spread(region,dur) %>%
  ungroup() %>% 
  as.data.frame()

list_regions <- names(aa_matrix[2:46]) 
aa_matrix[is.na(aa_matrix)] <- 0

response_matrix <- aa_corrected %>% 
  select(donor_id, SBS22, SBS1536I_cosmic, sex, age_diag) %>% 
  mutate(age_diag = as.numeric(age_diag))

all_matrix <- response_matrix %>% left_join(aa_matrix) %>% 
  mutate(across(!donor_id & !age_diag & !SBS22 & !SBS1536I_cosmic & !sex, ~./age_diag)) %>% 
  mutate(AA_sigs=SBS22+SBS1536I_cosmic) %>% 
  mutate(across(
    any_of(c("AA_sigs","SBS22","SBS1536I_cosmic")),
    .fns = list(int = ~qnorm((rank(.,na.last="keep")-0.5)/sum(!is.na(.))), pa = ~ifelse(.>0,1,0)),
    .names = "{.col}_{.fn}"))

# Mixed model - With data from AA
W <- all_matrix %>%
  select(-age_diag, -sex, -starts_with("SBS"), -starts_with("AA")) %>% 
  column_to_rownames(var="donor_id") %>% 
  as.matrix()

image(Matrix(t(W)), sub="", xlab="Subject", ylab="Region", oldstyle=F,
      scales=list(at=1:123,alternating=2,y=list(labels=colnames(W))))

M <- data.frame(select(all_matrix, -list_regions)) %>% mutate(region = rep(list_regions, length.out = NROW(W)))

# SBS22
lmod <- lFormula(SBS22_int ~ age_diag + sex + (1|region), data=M)
lmod$reTrms$Zt <- lmod$reTrms$Ztlist[[1]] <- Matrix(t(W))
devfun <- do.call(mkLmerDevfun, lmod)
opt <- optimizeLmer(devfun)
mSBS22 <- mkMerMod(environment(devfun), opt, lmod$reTrms, fr = lmod$fr)
plot(mSBS22)

ddSBS22 <- tidy(mSBS22, effects="ran_vals")
ddSBS22 <- transform(ddSBS22, level=reorder(level,estimate))
pSBS22 <- ggplot(ddSBS22,aes(x=level,y=estimate))+
  geom_pointrange(aes(ymin=estimate-1.96*std.error,
                      ymax=estimate+1.96*std.error))+coord_flip()
pSBS22

# SBS1536I_cosmic
lmod <- lFormula(SBS1536I_cosmic_int ~ age_diag + sex + (1|region), data=M)
lmod$reTrms$Zt <- lmod$reTrms$Ztlist[[1]] <- Matrix(t(W))
devfun <- do.call(mkLmerDevfun, lmod)
opt <- optimizeLmer(devfun)
mSBS1536I_cosmic <- mkMerMod(environment(devfun), opt, lmod$reTrms, fr = lmod$fr)
plot(mSBS1536I_cosmic)

ddSBS1536I_cosmic <- tidy(mSBS1536I_cosmic, effects="ran_vals")
ddSBS1536I_cosmic <- transform(ddSBS1536I_cosmic, level=reorder(level,estimate))
pSBS1536I_cosmic <- ggplot(ddSBS1536I_cosmic,aes(x=level,y=estimate))+
  geom_pointrange(aes(ymin=estimate-1.96*std.error,
                      ymax=estimate+1.96*std.error))+coord_flip()
pSBS1536I_cosmic

# Sum AA sigs
lmod <- lFormula(AA_sigs_int ~ age_diag + sex + (1|region), data=M)
lmod$reTrms$Zt <- lmod$reTrms$Ztlist[[1]] <- Matrix(t(W))
devfun <- do.call(mkLmerDevfun, lmod)
opt <- optimizeLmer(devfun)
mAA <- mkMerMod(environment(devfun), opt, lmod$reTrms, fr = lmod$fr)
plot(mAA)

ddAA <- tidy(mAA, effects="ran_vals")
ddAA <- transform(ddAA, level=reorder(level,estimate))
pAA <- ggplot(ddAA,aes(x=level,y=estimate))+
  geom_pointrange(aes(ymin=estimate-1.96*std.error,
                      ymax=estimate+1.96*std.error))+coord_flip() + theme_minimal()
pAA

# Plotting all estimates side by side
pSBS22+pSBS1536I_cosmic+pAA

estim_1536D <- select(ddSBS22, level, estimate, std.error) %>% 
  transmute(region = str_remove(level, "RO:|SR:"), SBS22_int=estimate, signif_1536D = ifelse(estimate + 1.96*std.error < 0 |estimate-1.96*std.error > 0, 1, 0))
estim_1536I <- select(ddSBS1536I_cosmic, level, estimate, std.error) %>% 
  transmute(region = str_remove(level, "RO:|SR:"), SBS1536I_cosmic_int=estimate, signif_1536I = ifelse(estimate + 1.96*std.error < 0 |estimate-1.96*std.error > 0, 1, 0))
estim_AA <- select(ddAA, level, estimate, std.error) %>% 
  transmute(region = str_remove(level, "RO:|SR:"), AA_sigs_int=estimate, signif_AA = ifelse(estimate + 1.96*std.error < 0 |estimate-1.96*std.error > 0, 1, 0))
estim_region <- estim_1536D %>% left_join(estim_1536I) %>% left_join(estim_AA)
rm(estim_1536D,estim_1536I,estim_AA)
saveRDS(estim_region, "output/estim_region")
